Fast CAGRA Index Merge by landrumb · Pull Request #2352 · NVIDIA/cuvs

landrumb · 2026-07-22T21:16:58Z

This PR implements the Fastener graph merge operation. This PR supports merging for float, half, int8, and uint8 dtypes, and euclidean distances.

There was originally specialization to use int8 GEMM for the integer types, but I ran into portability issues on Ada and it turns out that using the same f32 path for both is simpler and not substantially slower. Currently investigating switching the unified path to TF32 to use tensor cores.

The core logic resides in cagra_merge_scaffold.cuh.

This PR adds 7.64 MiB to libcuvs.so, a 2.94% increase.

(H100)

copy-pr-bot · 2026-07-22T21:17:01Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

copy-pr-bot · 2026-07-22T21:29:09Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

dantegd · 2026-07-22T21:32:37Z

/ok to test 2c54392

landrumb · 2026-07-23T19:24:32Z

/ok to test

landrumb · 2026-07-24T00:31:03Z

/ok to test

landrumb · 2026-07-24T19:09:05Z

/ok to test

landrumb · 2026-07-27T16:00:43Z

/ok to test

landrumb · 2026-07-30T21:34:31Z

Currently failing because the wheel builds are too big. Not sure there's much I can do about this.

achirkin

Thanks for this extensive contribution! Since the datataset API changes are in, the main request is to rethink where we allocate the new, consolidated dataset during merge (see below).
Then, please try to refactor the code to use raft mdarrays and mdspans in place of RMM and use the workspace/large_workspace resources for all temporary allocations - this will help us to keep track of memory requirements (which seem to be plenty).

achirkin · 2026-08-03T11:18:43Z

+/** Device state and tuning knobs shared by every split level: the precomputed row norms, the
+ * tiling and workspace capacities, and the seed feeding the deterministic leader samples. */
+struct split_context {
+  split_context(rmm::cuda_stream_view stream, int64_t rows)
+    : norms(static_cast<size_t>(rows), stream)
+  {
+  }
+
+  rmm::device_uvector<float> norms;
+  int assignment_tile_rows    = ASSIGNMENT_TILE_ROWS;
+  size_t gemm_workspace_bytes = GEMM_WORKSPACE_BYTES;
+  uint64_t seed               = DETERMINISTIC_SEED;
+};


Let's try to switch to raft primitives here: raft::device_vector<float> as the buffer and raft::resources as an argument.
Then, we need to consider: does this context fit into the bounded workspace memory or does it grow with the problem size unbounded? Depending on the answer, use raft::get_workspace_resource_ref or raft::get_large_workspace_resource_ref.
By doing this, we let the user control where the working memory for the algorithm resides, enable better accounting (via tools like raft::memory_tracking_resources), and prepare for the dry run execution (feature soon to be enabled in raft).

I've gone through and I think this is done correctly now, let me know if I misunderstand the usage of the raft allocs. Is there a way I should be testing preparedness for dry run execution?

achirkin · 2026-08-03T12:57:48Z

+  if (params.attach_dataset_on_build) {
+    // CAGRA search assumes 16-byte row alignment (vectorized loads). Consolidate remains dense for
+    // the scaffold; attach through make_aligned_dataset so unaligned dims are padded.
+    merged_index.update_dataset(handle, make_aligned_dataset(handle, std::move(dataset), 16));
+  } else {
+    using ds_idx_type = typename index<T, IdxT>::dataset_index_type;
+    merged_index.update_dataset(
+      handle, std::make_unique<cuvs::neighbors::empty_dataset<ds_idx_type>>(preflight.dim));
+  }


I guess, we should move the dataset consolidation into a separate place, perhaps as a part of the dataset API. CC @HowardHuang1 to discuss what would be the place for dataset merging under the new API and who should own it?

I can make an issue for this, but could we leave it here for now?

Adding offline context to the thread:
The dataset merging logic now present is just copying vectors into storage that's now provided by the caller per interface changes that were introduced with the dataset API. The rebuild version of this takes a bitmap to filter them, and the fastener version doesn't.

Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>

The graph merge PR (#2352) uses some of the kernels defined by CAGRA, and binary size takes a hit from them not appearing in a TU where they can be shared by both CAGRA and Fastener. Authors: - Ben Landrum (https://github.com/landrumb) Approvers: - Divye Gala (https://github.com/divyegala) URL: #2412

dantegd

Have only the small concerns, besides it things look solid!

Also, it looks like a few working artifacts may have slipped into the PR that need to be removed before merge: PR_2352_REVIEW_WORK.md, pipnn.pdf and FASTENER_PR_PLAN.md, right?

dantegd · 2026-08-06T14:40:27Z

+ * `cuvs::neighbors::cagra::detail::merged_dataset_size`.
+ * @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by


can users actually use this sizing instruction from the installed C++ API?

merged_dataset_size() is defined under cpp/src and isn’t declared in the public header. Could we either expose an appropriate public helper?

Looking more at it I'm not super convinced this is a useful enough helper to expose. It just gives the number of rows in the output, which is either a popcount of the filter or the sum of the inputs, and leaves the user to extract the dimension and stride to figure out the size of the alloc. I think it makes more sense to either surface a helper that gives the size and/or shape of the allocation or surface no helper at all.

julianmi · 2026-08-07T09:05:04Z

+struct merge_params {
+  merge_algo algo        = merge_algo::AUTO;
+  uint32_t levels        = 2;
+  uint32_t root_fanout   = 2;
+  uint32_t lower_fanout  = 3;
+  double leader_fraction = 0.02;
+  uint32_t max_leaders   = 1024;
+  uint32_t leaf_size     = 256;
+  uint32_t leaf_degree   = 4;
+};


We might want to add a section in the documentation on these parameters and their typical values.

landrumb added 7 commits July 16, 2026 18:15

factoring out in-place edge sort

637e0fc

core merge logic

7972d3a

remove VMM dataset consolidation from fastener merge

bc6ec7b

pedantic docstring change

e2ce96f

assert/validation changes

a9149db

test improvements

cf78e32

assert/message changes & aligned signatures

209485c

github-project-automation Bot added this to Unstructured Data Processing Jul 22, 2026

landrumb and others added 2 commits July 22, 2026 21:26

pre-commit hooks

da96f26

Merge branch 'main' into landrumb/cagra-fastener-merge

2c54392

aamijar assigned landrumb Jul 22, 2026

aamijar added non-breaking Introduces a non-breaking change feature request New feature or request labels Jul 22, 2026

aamijar added this to Unstructured Data Processing Jul 22, 2026

aamijar moved this to In Progress in Unstructured Data Processing Jul 22, 2026

landrumb added 2 commits July 22, 2026 22:16

clearer check rationale

9a3cabf

test arg mismatch

af25e86

landrumb and others added 3 commits July 23, 2026 21:05

refactor to use existing select_k kernel

2eb64fa

fixed attached dataset alignment issue

6f65f87

Merge branch 'main' into landrumb/cagra-fastener-merge

2e6e148

CI fixes

c8fd540

aligned test with earlier ci fix

951b0e2

Merge branch 'main' into landrumb/cagra-fastener-merge

678f2ea

landrumb added 3 commits July 30, 2026 17:21

missed files on default increase

298bfcf

removed obsolete template

8891aae

style miss

38c59f7

achirkin requested changes Aug 3, 2026

View reviewed changes

landrumb and others added 5 commits August 3, 2026 10:23

artem's suggestions

9dfa693

Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>

additional suggestion

5a74ddb

Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>

merging main

1cba27d

replacing rmm + non-constant workspace

87fafd3

signature mismatch

879e6b0

landrumb requested a review from dantegd August 4, 2026 16:44

landrumb added 2 commits August 4, 2026 17:12

switching to cublasltmatmul

0c3f8a4

missed an alloc

653e8a0

landrumb force-pushed the landrumb/cagra-fastener-merge branch from 1b8558a to 653e8a0 Compare August 4, 2026 20:20

temporary factoring out of shared kernels

cd83283

landrumb mentioned this pull request Aug 4, 2026

Factor CAGRA graph kernels into shared translation unit #2412

Merged

landrumb requested a review from achirkin August 5, 2026 17:50

landrumb mentioned this pull request Aug 5, 2026

[FEA] Supporting filters in the Fastener implementation of CAGRA merge #2417

Open

dantegd requested changes Aug 6, 2026

View reviewed changes

landrumb added 4 commits August 6, 2026 16:44

deleting bycatch

6b19de2

merging main (w/ CAGRA TU PR)

df77ae9

smarter cast

a7e294b

moved leaf size check and added test

7e360dc

julianmi reviewed Aug 7, 2026

View reviewed changes

landrumb added 5 commits August 7, 2026 20:38

replacing double with float

4463256

one warp per row copy

aa2c143

writing scaffold directly to output graph

a2b378c

selective zeroing of the output

a5f7077

aligning C failure with C++

171be76

		* `cuvs::neighbors::cagra::detail::merged_dataset_size`.
		* @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by

Uh oh!

Conversation

landrumb commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

dantegd commented Jul 22, 2026

Uh oh!

landrumb commented Jul 23, 2026

Uh oh!

landrumb commented Jul 24, 2026

Uh oh!

landrumb commented Jul 24, 2026

Uh oh!

landrumb commented Jul 27, 2026

Uh oh!

landrumb commented Jul 30, 2026

Uh oh!

achirkin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dantegd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

landrumb commented Jul 22, 2026 •

edited

Loading